#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
图片压缩脚本

功能：
- 压缩指定目录下的所有图片文件
- 将压缩后的图片输出到另一个目录
- 压缩规则：保持图片比例，最长边小于指定像素值
- 支持常见图片格式（jpg, jpeg, png, bmp, gif等）
- 提供命令行参数：输入目录、输出目录、最大像素值
- 添加进度显示和错误处理
- 保持原始文件名和目录结构

使用方法：
python image_compressor.py -i /path/to/input -o /path/to/output -s 800

作者：AI Assistant
日期：2024
"""

import os
import sys
import argparse
from pathlib import Path
from typing import List, Tuple
import logging
from concurrent.futures import ThreadPoolExecutor, as_completed
import threading
from PIL import Image, ImageOps
from tqdm import tqdm


# 支持的图片格式
SUPPORTED_FORMATS = {'.jpg', '.jpeg', '.png', '.bmp', '.gif', '.tiff', '.webp'}


def setup_logging(log_level: str = 'INFO') -> None:
    """
    设置日志配置
    
    Args:
        log_level: 日志级别
    """
    logging.basicConfig(
        level=getattr(logging, log_level.upper()),
        format='%(asctime)s - %(levelname)s - %(message)s',
        handlers=[
            logging.StreamHandler(sys.stdout),
            logging.FileHandler('image_compressor.log', encoding='utf-8')
        ]
    )


def get_image_files(input_dir: Path) -> List[Path]:
    """
    获取目录下所有支持的图片文件
    
    Args:
        input_dir: 输入目录路径
        
    Returns:
        图片文件路径列表
    """
    image_files = []
    
    for file_path in input_dir.rglob('*'):
        if file_path.is_file() and file_path.suffix.lower() in SUPPORTED_FORMATS:
            image_files.append(file_path)
    
    return image_files


def calculate_new_size(original_size: Tuple[int, int], max_size: int) -> Tuple[int, int]:
    """
    计算新的图片尺寸，保持比例
    
    Args:
        original_size: 原始尺寸 (width, height)
        max_size: 最长边的最大像素值
        
    Returns:
        新的尺寸 (width, height)
    """
    width, height = original_size
    
    # 如果图片已经小于等于指定尺寸，不需要压缩
    if max(width, height) <= max_size:
        return original_size
    
    # 计算缩放比例
    if width > height:
        ratio = max_size / width
    else:
        ratio = max_size / height
    
    new_width = int(width * ratio)
    new_height = int(height * ratio)
    
    return new_width, new_height


def compress_image(input_path: Path, output_path: Path, max_size: int, quality: int = 85) -> bool:
    """
    压缩单个图片文件
    
    Args:
        input_path: 输入图片路径
        output_path: 输出图片路径
        max_size: 最长边的最大像素值
        quality: JPEG质量（1-100）
        
    Returns:
        是否成功压缩
    """
    try:
        # 确保输出目录存在
        output_path.parent.mkdir(parents=True, exist_ok=True)
        
        # 打开图片
        with Image.open(input_path) as img:
            # 自动旋转图片（处理EXIF方向信息）
            img = ImageOps.exif_transpose(img)
            
            # 计算新尺寸
            new_size = calculate_new_size(img.size, max_size)
            
            # 如果需要压缩
            if new_size != img.size:
                # 使用高质量的重采样算法
                img = img.resize(new_size, Image.Resampling.LANCZOS)
            
            # 保存图片
            save_kwargs = {}
            if output_path.suffix.lower() in {'.jpg', '.jpeg'}:
                # JPEG格式设置质量和优化
                save_kwargs.update({
                    'quality': quality,
                    'optimize': True,
                    'progressive': True
                })
                # 如果原图有透明通道，转换为RGB
                if img.mode in ('RGBA', 'LA', 'P'):
                    background = Image.new('RGB', img.size, (255, 255, 255))
                    if img.mode == 'P':
                        img = img.convert('RGBA')
                    background.paste(img, mask=img.split()[-1] if img.mode == 'RGBA' else None)
                    img = background
            elif output_path.suffix.lower() == '.png':
                # PNG格式优化
                save_kwargs.update({
                    'optimize': True,
                    'compress_level': 6
                })
            
            img.save(output_path, **save_kwargs)
            
        return True
        
    except Exception as e:
        logging.error(f"压缩图片失败 {input_path}: {str(e)}")
        return False


def process_single_image(args_tuple: Tuple[Path, Path, Path, int, int]) -> Tuple[bool, int, int]:
    """
    处理单个图片的包装函数，用于多线程调用
    
    Args:
        args_tuple: (input_path, input_dir, output_dir, max_size, quality)
        
    Returns:
        (是否成功, 原始大小, 压缩后大小)
    """
    input_path, input_dir, output_dir, max_size, quality = args_tuple
    
    try:
        # 计算相对路径，保持目录结构
        relative_path = input_path.relative_to(input_dir)
        output_path = output_dir / relative_path
        
        # 记录原始文件大小
        original_size = input_path.stat().st_size
        
        # 压缩图片
        if compress_image(input_path, output_path, max_size, quality):
            # 记录压缩后文件大小
            compressed_size = 0
            if output_path.exists():
                compressed_size = output_path.stat().st_size
                
                # 计算压缩比例
                compression_ratio = (1 - compressed_size / original_size) * 100
                logging.debug(f"压缩完成: {relative_path} (压缩率: {compression_ratio:.1f}%)")
            
            return True, original_size, compressed_size
        else:
            return False, original_size, 0
            
    except Exception as e:
        logging.error(f"处理文件时出错 {input_path}: {str(e)}")
        return False, 0, 0


def compress_images(input_dir: Path, output_dir: Path, max_size: int, quality: int = 85, max_workers: int = None) -> None:
    """
    批量压缩图片（多线程版本）
    
    Args:
        input_dir: 输入目录
        output_dir: 输出目录
        max_size: 最长边的最大像素值
        quality: JPEG质量（1-100）
        max_workers: 最大线程数，默认为CPU核心数
    """
    # 获取所有图片文件
    image_files = get_image_files(input_dir)
    
    if not image_files:
        logging.warning(f"在目录 {input_dir} 中没有找到支持的图片文件")
        return
    
    # 设置默认线程数
    if max_workers is None:
        max_workers = min(32, (os.cpu_count() or 1) + 4)  # 限制最大线程数，避免过度占用资源
    
    logging.info(f"找到 {len(image_files)} 个图片文件")
    logging.info(f"使用 {max_workers} 个线程进行并行处理")
    
    # 统计信息（使用锁保证线程安全）
    stats_lock = threading.Lock()
    success_count = 0
    error_count = 0
    total_original_size = 0
    total_compressed_size = 0
    
    # 准备任务参数
    tasks = [(img_path, input_dir, output_dir, max_size, quality) for img_path in image_files]
    
    # 使用线程池执行任务
    with ThreadPoolExecutor(max_workers=max_workers) as executor:
        # 提交所有任务
        future_to_task = {executor.submit(process_single_image, task): task for task in tasks}
        
        # 使用进度条显示处理进度
        with tqdm(total=len(image_files), desc="压缩图片", unit="张") as pbar:
            for future in as_completed(future_to_task):
                try:
                    success, original_size, compressed_size = future.result()
                    
                    # 线程安全地更新统计信息
                    with stats_lock:
                        if success:
                            success_count += 1
                            total_original_size += original_size
                            total_compressed_size += compressed_size
                        else:
                            error_count += 1
                            
                except Exception as e:
                    with stats_lock:
                        error_count += 1
                    logging.error(f"任务执行异常: {str(e)}")
                
                # 更新进度条
                pbar.update(1)
    
    # 输出统计信息
    logging.info(f"压缩完成！")
    logging.info(f"成功: {success_count} 张")
    logging.info(f"失败: {error_count} 张")
    
    if total_original_size > 0:
        total_compression_ratio = (1 - total_compressed_size / total_original_size) * 100
        logging.info(f"原始总大小: {total_original_size / 1024 / 1024:.2f} MB")
        logging.info(f"压缩后总大小: {total_compressed_size / 1024 / 1024:.2f} MB")
        logging.info(f"总体压缩率: {total_compression_ratio:.1f}%")


def validate_args(args) -> bool:
    """
    验证命令行参数
    
    Args:
        args: 解析后的命令行参数
        
    Returns:
        参数是否有效
    """
    # 检查输入目录
    if not args.input.exists():
        logging.error(f"输入目录不存在: {args.input}")
        return False
    
    if not args.input.is_dir():
        logging.error(f"输入路径不是目录: {args.input}")
        return False
    
    # 检查最大尺寸
    if args.size <= 0:
        logging.error(f"最大尺寸必须大于0: {args.size}")
        return False
    
    # 检查质量参数
    if not (1 <= args.quality <= 100):
        logging.error(f"质量参数必须在1-100之间: {args.quality}")
        return False
    
    return True


def main():
    inputDir="/Users/wujinjie/Documents/work/data/cdn/img"
    outputDir="/Users/wujinjie/Documents/work/data/cdn/mini500"
    
    # 设置日志
    setup_logging()
    
    try:
        # 开始压缩（使用多线程，默认线程数为CPU核心数）
        compress_images(Path(inputDir), Path(outputDir), 500, 85, max_workers=8)
        
    except KeyboardInterrupt:
        logging.info("用户中断操作")
        sys.exit(1)
    except Exception as e:
        logging.error(f"程序执行出错: {str(e)}")
        sys.exit(1)


if __name__ == '__main__':
    main()